import logging
from django.db import connections

from yida_app_service.utils.custom_db_utils import AbstractDBTool, CustomDataSourceManager

class DBTool(AbstractDBTool):
    def __init__(self, database,anti_sql_injection="%s"):
        super().__init__(anti_sql_injection=anti_sql_injection)
        """
        初始化DBTool类，并指定数据库。

        :param database: str，连接的数据库名称。如果未提供数据库名称，将抛出异常。
        """
        if not database:
            error_message = "数据库名称是必填项，不能为None或空字符串。"
            logging.error(error_message)
            raise ValueError(error_message)
        self.database = database

    def execute_query(self, query, params=None,return_dict=False):
        """
        在指定数据库上执行原始查询。

        :param query: str，要执行的SQL查询。
        :param params: tuple或list，包含在查询中的参数。
        :return: list，返回查询到的结果行。
        """
        try:
            with connections[self.database].cursor() as cursor:
                cursor.execute(query, params or [])
                rows = cursor.fetchall()
                # 如果要返回字典
                if return_dict:
                    return CustomDataSourceManager.to_dict(cursor=cursor,rows=rows)
                return rows
        except Exception as e:
            logging.error(f"执行查询时发生错误: {str(e)}")
            return []

    def execute_update(self, query, params=None):
        """
        在指定数据库上执行UPDATE（或INSERT/DELETE）查询。

        :param query: str，要执行的SQL查询。
        :param params: tuple或list，包含在查询中的参数。
        :return: int，受影响的行数，出错时返回-1。
        """
        try:
            with connections[self.database].cursor() as cursor:
                cursor.execute(query, params or [])
                return cursor.rowcount  # 返回受影响的行数
        except Exception as e:
            logging.error(f"执行更新时发生错误: {str(e)}")
            return -1
